#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RK3588端PLC控制客户端
严格按照割叶机通讯协议说明实现
"""

import json
import logging
import socket
import threading
import time
import sys
import random
import os

# 兼容不同版本的pymodbus
try:
    from pymodbus.client import ModbusSerialClient
except ImportError:
    try:
        from pymodbus.client.sync import ModbusSerialClient
    except ImportError:
        from pymodbus.client.serial import ModbusSerialClient

try:
    from pymodbus.exceptions import ModbusException
except ImportError:
    class ModbusException(Exception):
        pass

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s'
)
logger = logging.getLogger(__name__)

class PLCController:
    """PLC控制器 - 严格按照割叶机通讯协议说明"""
    
    def __init__(self, port='/dev/ttyUSB0'):
        # 严格按照协议说明的参数
        self.port = port
        self.baudrate = 9600        # 协议规定：9600bit/s
        self.slave_id = 1           # 协议规定：通讯地址1
        self.bytesize = 8           # 协议规定：8位数据位
        self.parity = 'N'           # 协议规定：无奇偶校验
        self.stopbits = 1           # 协议规定：1位停止位
        
        self.client = None
        self.connected = False
        
        # 按照协议说明的寄存器映射
        self.registers = {
            # 控制寄存器 (功能码06/10)
            'control_authority': 120,    # D120 上位机控制权限
            'turn_control': 121,         # D121 转向控制 (-12到12)
            'move_control': 122,         # D122 行进控制 (-25到25)
            'generator_control': 123,    # D123 发电机启动控制
            'relay_220v': 124,          # D124 220V继电器控制
            'water_spray': 125,         # D125 喷水控制
            'cooling_control': 126,     # D126 制冷控制
            'speed_percent': 127,       # D127 行进速度百分比
            'leaf_cutter': 128,         # D128 割叶机控制
            
            # 传感器寄存器 (功能码03)
            'generator_voltage': 129,    # D129 发电机电压
            'generator_current': 130,    # D130 发电机电流
            'generator_power': 131,      # D131 发电机功率
            'generator_temp': 132,       # D132 发电机温度
            'engine_temp': 133,          # D133 发动机温度
            'hot_water_temp': 134,       # D134 热水端温度
            'cooling_temp': 135,         # D135 制冷端温度
            'box_temp': 136,             # D136 箱体温度
            'fuel_level': 137,           # D137 油箱液位
            'water_level': 138,          # D138 水箱液位
            'plc_output_status': 139     # D139 PLC输出状态
        }
        
        logger.info("📋 PLC控制器初始化完成")
        logger.info(f"   串口: {self.port}")
        logger.info(f"   波特率: {self.baudrate} bit/s")
        logger.info(f"   通讯地址: {self.slave_id}")
        logger.info(f"   数据位: {self.bytesize}位")
        logger.info(f"   校验位: {self.parity}")
        logger.info(f"   停止位: {self.stopbits}位")
    
    def connect(self):
        """连接PLC - 严格按照协议参数"""
        logger.info("🔌 开始连接PLC设备...")
        
        # 检查串口设备
        if not os.path.exists(self.port):
            logger.error(f"❌ 串口设备不存在: {self.port}")
            logger.error("   请检查USB转485设备连接")
            return False
        
        # 尝试连接
        for attempt in range(3):
            if attempt > 0:
                logger.info(f"🔄 连接重试 {attempt+1}/3...")
                time.sleep(1)
            
            try:
                # 创建Modbus客户端
                try:
                    # 新版本pymodbus
                    self.client = ModbusSerialClient(
                        port=self.port,
                        baudrate=self.baudrate,
                        bytesize=self.bytesize,
                        parity=self.parity,
                        stopbits=self.stopbits,
                        timeout=2
                    )
                except TypeError:
                    # 旧版本pymodbus
                    self.client = ModbusSerialClient(
                        method='rtu',
                        port=self.port,
                        baudrate=self.baudrate,
                        bytesize=self.bytesize,
                        parity=self.parity,
                        stopbits=self.stopbits,
                        timeout=2
                    )
                
                # 尝试连接
                if self.client.connect():
                    logger.info("✅ 串口连接成功")
                    
                    # 测试读取寄存器D120
                    logger.info("🧪 测试PLC通信...")
                    try:
                        result = self._test_read_register(121)
                        if result is not None:
                            logger.info(f"🎉 PLC通信成功！寄存器D120值: {result}")
                            self.connected = True
                            
                            # 不自动设置控制权限，保持PLC原始状态
                            logger.info("✅ PLC连接成功，保持原始控制模式")
                            
                            return True
                        else:
                            logger.warning(f"❌ PLC无响应 (尝试 {attempt+1}/3)")
                    except Exception as e:
                        logger.warning(f"❌ PLC通信测试失败: {e}")
                    
                    self.client.close()
                else:
                    logger.warning(f"❌ 串口连接失败 (尝试 {attempt+1}/3)")
                    
            except Exception as e:
                logger.error(f"❌ 连接异常 (尝试 {attempt+1}/3): {e}")
        
        logger.error("❌ PLC连接失败")
        logger.error("🔧 请检查:")
        logger.error("   1. PLC设备是否通电")
        logger.error("   2. 485接线: 黄色→A, 白色→B")
        logger.error("   3. PLC通讯地址是否为1")
        logger.error("   4. PLC波特率是否为9600")
        return False
    
    def _test_read_register(self, address):
        """测试读取寄存器"""
        try:
            # 兼容不同版本API
            try:
                result = self.client.read_holding_registers(address=address, count=1, slave=self.slave_id)
            except TypeError:
                try:
                    result = self.client.read_holding_registers(address, 1, slave=self.slave_id)
                except TypeError:
                    result = self.client.read_holding_registers(address, 1, unit=self.slave_id)
            
            # 检查结果
            if hasattr(result, 'isError') and callable(getattr(result, 'isError')):
                if not result.isError():
                    return result.registers[0]
            elif hasattr(result, 'registers'):
                return result.registers[0]
            
            return None
        except Exception:
            return None
    
    def disconnect(self):
        """断开PLC连接"""
        if self.client:
            self.client.close()
            self.connected = False
            logger.info("Modbus连接已断开")
    
    
    def read_register(self, register_name):
        """读取寄存器"""
        if not self.connected or register_name not in self.registers:
            return None
        
        try:
            address = self.registers[register_name]
            result = self._test_read_register(address)
            
            if result is not None:
                # 处理有符号数据
                if result > 32767:
                    result = result - 65536
                return result
            else:
                logger.warning(f"读取寄存器{register_name}失败")
                return None
                
        except Exception as e:
            logger.error(f"读取寄存器{register_name}异常: {e}")
            return None
    
    def write_register(self, register_name, value):
        """写入寄存器"""
        if not self.connected or register_name not in self.registers:
            return False
        
        # 安全检查：控制权限寄存器需要特殊处理
        if register_name == 'control_authority':
            logger.warning(f"⚠️ 控制权限变更请求: {value} (0=本地遥控, 1=远程控制)")
            logger.warning("   这可能影响伺服电机和原生遥控器功能")
            logger.info("🔒 为安全起见，禁用控制权限变更")
            return False
        
        try:
            address = self.registers[register_name]
            
            # 处理有符号数据
            if value < 0:
                value = value + 65536
            
            # 兼容不同版本API
            try:
                result = self.client.write_register(address=address, value=value, slave=self.slave_id)
            except TypeError:
                try:
                    result = self.client.write_register(address, value, slave=self.slave_id)
                except TypeError:
                    result = self.client.write_register(address, value, unit=self.slave_id)
            
            # 检查结果
            if hasattr(result, 'isError') and callable(getattr(result, 'isError')):
                if not result.isError():
                    logger.info(f"写入寄存器{register_name}={value}")
                    return True
                else:
                    logger.warning(f"写入寄存器{register_name}失败")
                    return False
            else:
                logger.info(f"写入寄存器{register_name}={value}")
                return True
                
        except Exception as e:
            logger.error(f"写入寄存器{register_name}异常: {e}")
            return False
    
    def read_all_sensors(self):
        """读取所有传感器数据"""
        sensors = {}
        
        # 传感器寄存器列表
        sensor_registers = [
            'generator_voltage', 'generator_current', 'generator_power',
            'generator_temp', 'engine_temp', 'hot_water_temp',
            'cooling_temp', 'box_temp', 'fuel_level', 'water_level',
            'plc_output_status'
        ]
        
        for register in sensor_registers:
            value = self.read_register(register)
            if value is not None:
                sensors[register] = value
        
        return sensors

class RK3588Client:
    """RK3588客户端主类"""
    
    def __init__(self, server_host, server_port):
        self.server_host = server_host
        self.server_port = server_port
        self.socket = None
        self.connected = False
        self.running = True
        
        # 初始化PLC控制器
        self.plc = PLCController()
        
        # 线程
        self.heartbeat_thread = None
        self.command_thread = None
    
    def connect_to_server(self):
        """连接到服务器"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            
            # 设置socket选项防止TIME_WAIT问题
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            self.socket.settimeout(10)  # 设置10秒连接超时（与旧代码一致）
            
            logger.info(f"正在连接服务器: {self.server_host}:{self.server_port}")
            self.socket.connect((self.server_host, self.server_port))
            
            self.connected = True
            logger.info(f"✅ 连接到服务器成功: {self.server_host}:{self.server_port}")
            
            # 发送初始状态
            self.send_status()
            return True
            
        except socket.timeout:
            logger.error(f"❌ 连接服务器超时: {self.server_host}:{self.server_port}")
            logger.error("   请检查:")
            logger.error("   1. 服务器IP地址是否正确")
            logger.error("   2. 服务器是否已启动 (python app.py)")
            logger.error("   3. 网络连接是否正常 (ping测试)")
            logger.error("   4. 防火墙是否允许9000端口")
            return False
        except ConnectionRefusedError:
            logger.error(f"❌ 连接被拒绝: {self.server_host}:{self.server_port}")
            logger.error("   服务器可能未启动或端口未监听")
            return False
        except Exception as e:
            logger.error(f"❌ 连接服务器失败: {e}")
            logger.error(f"   目标: {self.server_host}:{self.server_port}")
            return False
    
    def disconnect_from_server(self):
        """断开服务器连接"""
        if self.socket:
            try:
                # 优雅关闭：先关闭发送，再关闭接收
                self.socket.shutdown(socket.SHUT_RDWR)
            except:
                pass  # 可能已经关闭
            finally:
                self.socket.close()
                self.socket = None
                self.connected = False
                logger.info("服务器连接已断开")
    
    def send_message(self, message):
        """发送消息到服务器"""
        if not self.connected or not self.socket:
            return False
        
        try:
            data = json.dumps(message, ensure_ascii=False) + '\n'
            self.socket.send(data.encode('utf-8'))
            return True
        except Exception as e:
            logger.error(f"发送消息失败: {e}")
            self.connected = False
            return False
    
    def send_status(self):
        """发送状态信息 - 只发送连接状态，不主动读取传感器"""
        status = {
            'type': 'status',
            'timestamp': time.time(),
            'plc_connected': self.plc.connected
        }
        self.send_message(status)
    
    def receive_commands(self):
        """接收服务器命令"""
        while self.running:
            try:
                if not self.connected:
                    time.sleep(1)
                    continue
                
                self.socket.settimeout(0.1)
                data = self.socket.recv(1024)
                if not data:
                    break
                
                message = data.decode('utf-8').strip()
                if not message:
                    continue
                
                # 处理多条JSON数据
                lines = message.split('\n')
                for line in lines:
                    line = line.strip()
                    if line and line.startswith('{'):
                        try:
                            command = json.loads(line)
                            self.process_command(command)
                        except json.JSONDecodeError as je:
                            logger.error(f"JSON解析失败: {je}")
                            continue
                
            except socket.timeout:
                continue
            except Exception as e:
                logger.error(f"接收命令失败: {e}")
                self.connected = False
                break
    
    def check_usb_device(self):
        """检查USB转485设备"""
        import glob
        
        # 检查常见的USB串口设备
        usb_devices = glob.glob('/dev/ttyUSB*')
        
        if usb_devices:
            return {
                'success': True,
                'devices': usb_devices,
                'selected': usb_devices[0] if usb_devices else None
            }
        else:
            return {
                'success': False,
                'message': 'USB转485设备未检测到'
            }
    
    def connect_plc_on_demand(self):
        """按需连接PLC"""
        if self.plc.connected:
            return {
                'success': True,
                'message': 'PLC已连接',
                'already_connected': True
            }
        
        # 尝试连接PLC
        success = self.plc.connect()
        
        return {
            'success': success,
            'message': 'PLC连接成功' if success else 'PLC连接失败',
            'plc_connected': self.plc.connected
        }
    
    def process_command(self, command):
        """处理服务器命令 - 只允许安全操作"""
        try:
            cmd_type = command.get('type')
            
            if cmd_type == 'check_usb':
                # 检查USB设备
                result = self.check_usb_device()
                response = {
                    'type': 'usb_check_result',
                    'result': result,
                    'timestamp': time.time()
                }
                self.send_message(response)
                logger.info(f"USB检测结果: {result}")
            
            elif cmd_type == 'connect_plc':
                # 连接PLC
                result = self.connect_plc_on_demand()
                response = {
                    'type': 'plc_connect_result',
                    'result': result,
                    'timestamp': time.time()
                }
                self.send_message(response)
                logger.info(f"PLC连接结果: {result}")
            
            elif cmd_type == 'movement':
                # 只允许移动控制：前进/后退/左转/右转
                forward = command.get('forward', 0)  # D122 行进控制 (-25到25)
                turn = command.get('turn', 0)        # D121 转向控制 (-12到12)
                
                # 限制范围
                forward = max(-25, min(25, forward))
                turn = max(-12, min(12, turn))
                
                # 写入移动寄存器
                success_move = self.plc.write_register('move_control', forward)
                success_turn = self.plc.write_register('turn_control', turn)
                
                response = {
                    'type': 'movement_response',
                    'success': success_move and success_turn,
                    'forward': forward,
                    'turn': turn,
                    'timestamp': time.time()
                }
                self.send_message(response)
                logger.info(f"移动控制: 前进={forward}, 转向={turn}")
            
            elif cmd_type == 'read_sensors':
                # 读取所有传感器数据
                sensors = self.plc.read_all_sensors()
                
                response = {
                    'type': 'sensor_data',
                    'sensors': sensors,
                    'plc_connected': self.plc.connected,
                    'timestamp': time.time()
                }
                self.send_message(response)
                logger.info("发送传感器数据")
            
            elif cmd_type == 'ping':
                # 处理ping命令
                response = {
                    'type': 'pong',
                    'timestamp': time.time()
                }
                self.send_message(response)
                
        except Exception as e:
            logger.error(f"处理命令失败: {e}")
    
    def heartbeat_loop(self):
        """心跳循环"""
        while self.running:
            try:
                if self.connected:
                    self.send_status()
                time.sleep(2)  # 每2秒发送一次状态
            except Exception as e:
                logger.error(f"心跳失败: {e}")
                time.sleep(5)
    
    def start_threads(self):
        """启动工作线程"""
        self.heartbeat_thread = threading.Thread(target=self.heartbeat_loop, daemon=True)
        self.command_thread = threading.Thread(target=self.receive_commands, daemon=True)
        
        self.heartbeat_thread.start()
        self.command_thread.start()
        
        logger.info("工作线程已启动")
    
    def run(self):
        """运行客户端"""
        logger.info("启动RK3588客户端: {}:{}".format(self.server_host, self.server_port))
        
        # 只连接服务器，不主动连接PLC
        logger.info("⏸️ 等待上位机触发PLC连接...")
        
        # 连接服务器
        if not self.connect_to_server():
            logger.error("❌ 服务器连接失败")
            return False
        
        # 启动工作线程
        self.start_threads()
        
        logger.info("✅ RK3588客户端启动完成")
        
        try:
            # 主循环
            while self.running:
                time.sleep(1)
                
                # 检查连接状态
                if not self.connected:
                    logger.info("尝试重新连接服务器...")
                    if self.connect_to_server():
                        logger.info("✅ 服务器重连成功")
                    else:
                        time.sleep(5)
                        
        except KeyboardInterrupt:
            logger.info("收到停止信号")
        finally:
            self.cleanup()
    
    def cleanup(self):
        """清理资源"""
        logger.info("正在清理资源...")
        self.running = False
        
        if self.plc and self.plc.connected:
            # 紧急恢复：确保退出时恢复本地控制
            logger.info("🚨 紧急恢复：确保本地遥控权限")
            try:
                # 临时允许写入控制权限寄存器进行恢复
                address = self.plc.registers['control_authority']
                try:
                    result = self.plc.client.write_register(address=address, value=0, slave=self.plc.slave_id)
                except TypeError:
                    try:
                        result = self.plc.client.write_register(address, 0, slave=self.plc.slave_id)
                    except TypeError:
                        result = self.plc.client.write_register(address, 0, unit=self.plc.slave_id)
                
                logger.info("✅ 本地遥控权限已恢复")
            except Exception as e:
                logger.error(f"❌ 恢复本地控制失败: {e}")
            
            self.plc.disconnect()
        
        if self.socket:
            self.disconnect_from_server()
        
        logger.info("资源清理完成")

def main():
    """主函数"""
    if len(sys.argv) != 3:
        print("用法: python3 rk3588_client.py <服务器IP> <端口>")
        sys.exit(1)
    
    server_host = sys.argv[1]
    server_port = int(sys.argv[2])
    
    client = RK3588Client(server_host, server_port)
    client.run()

if __name__ == "__main__":
    main()
